Use MaxText max_segments_per_seq config variable to control Grain batch packing #2774
Conversation
gabeweisz
requested review from
NicoGrande,
SurbhiJainUSC,
aireenmei,
richjames0 and
shralex
as code owners
yeandy
left a comment
yeandy
left a comment
There was a problem hiding this comment.
We may need to add max_sequences_per_bin=config.max_segments_per_seq in make_hf_eval_iterator too
Done, thanks for the tip |
gabeweisz
requested review from
A9isha,
NuojCheng,
RissyRan,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jiangjy1982,
khatwanimohit,
parambole,
shuningjin,
suexu1025 and
vipannalla
as code owners
|
Can you also update here: https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/configs/types.py#L882? And apply max_sequences_per_bin here: https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/input_pipeline/_grain_data_processing.py#L247 |
gabeweisz
requested review from
Obliviour,
SujeethJinesh,
mitalisi,
notabee and
shauryagup
as code owners
gabeweisz
requested review from
Lumosis,
gpolovets1,
jrplatin,
lukebaumann,
mailvijayasingh,
michelle-yooh,
patemotter and
xuefgu
as code owners
gabeweisz
force-pushed
the
gw_plumb_max_sequences_per_seq_to_grain
branch
from
9f3b716 to
ecae099
Compare
Done |
gabeweisz
force-pushed
the
gw_plumb_max_sequences_per_seq_to_grain
branch
from
1c5b56e to
d107901
Compare
|
@aireenmei I apologize but I needed to make some updates for the code quality checks (they didn't run before your approval). Can you please take another look. These were minor edits. |
|
Could you take a look to see if this error is related: https://github.com/AI-Hypercomputer/maxtext/actions/runs/20354898076/job/58493556179?pr=2774 |
|
I've asked @SurbhiJainUSC to also take a look. @gabeweisz you will need to squash your commits into 1 before merging |
gabeweisz
force-pushed
the
gw_plumb_max_sequences_per_seq_to_grain
branch
from
d107901 to
f7971f2
Compare
|
@aireenmei thanks for the feedback - I have done the squash commit |
|
Looks like I still need an approval from a maintainer. How can I get one? |
4 participants
Description
When using THD format packed data with TransformerEngine, the user must specify the maximum number of segments that can be packed into a sequence at Jax JIT time. If grain packs more segments than allowed, then this can cause crashes or data corruption.
We have previously updated grain to allow limiting the number of segments to pack into a sequence, and this PR takes the appropriate value from the MaxText configuration and passes it to Grain
Tests
We have had this fix in place in our AMD fork of MaxText for some time, but needed to get the Grain fix upstreamed first before creating this PR.
We have tested this fix extensively internally and have customers using it in production.
MaxText does not currently have any tests that use packed batches, but I can create some if needed.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.